from blackrltask import BlackRLTask
from blackrlenv  import BlackRLEnv

from pybrain.rl.learners.valuebased import ActionValueTable
from pybrain.rl.agents import LearningAgent
from pybrain.rl.learners import Q
from pybrain.rl.experiments import Experiment
from pybrain.rl.explorers import EpsilonGreedyExplorer

av_table = ActionValueTable(21, 2)
av_table.initialize(0.)

# define Q-learning agent
learner = Q(0.5, 0.0)
learner._setExplorer(EpsilonGreedyExplorer(0.0))
agent = LearningAgent(av_table, learner)

# define the environment
env = BlackRLEnv()

# define the task
task = BlackRLTask(env)

# finally, define experiment
experiment = Experiment(task, agent)

# ready to go, start the process
for x in range(0, 5):
    experiment.doInteractions(1)
    agent.learn()
    agent.reset()

q_matrix = av_table.params.reshape(21, 2)
print q_matrix
